import os
import argparse
import json
import math
import sys
import time

grounding = []
grounding_train = "grounding_train_20000.json"
grounding_valid = "grounding_valid_20000.json"
grounding_test = "grounding_test_20000.json"

grounding.append(grounding_train)
grounding.append(grounding_valid)
grounding.append(grounding_test)

def openreadtxt(file_name):
    data = []
    file = open(file_name, 'r')  # 打开文件
    file_data = file.readlines()  # 读取所有行
    for row in file_data:
        tmp_list = row.split(' ')  # 按‘ ’切分每行的数据
        tmp_list[-1] = tmp_list[-1].replace('\n','') #去掉换行符
        data.append(tmp_list)  # 将每行数据插入data中
    return data

def seek_files(image_dir, img_name):
    '''文件夹中查找文件'''
    for root, dirs, files in os.walk(image_dir):
        if img_name in files:
            # print('{} exists'.format(img_name))
            return True
        else:
            # print('{} not exists'.format(img_name))
            return False

def fit_json(image_dir, old_json, new_json):
    '''在图片目录中查找grounding文件中使用的图片，没有则将图片对应json信息写入new.json'''
    cnt = 0
    remain = 0
    dict = []

    for i in range(3):
        meta_json = os.path.join(old_json, grounding[i])
        print('checking {}'.format(meta_json))
        if not os.path.exists(meta_json):
            print('[ERROR] input_metadata_json does not exist.')

        meta_data = json.load(open(meta_json))
        total = len(meta_data)
        for item in meta_data:
            cnt += 1
            img_name = item['image']
            # 如果目录中有该照片，写入新文件
            # img_name = "VOA_EN_NW_2009.11.02.414022_0.jpg"
            # if seek_files(image_dir, img_name, dict):
            #     remain += 1
            #     dict.append(item)
            # 如果目录中没有该照片，写入新文件
            if not seek_files(image_dir, img_name):
                remain += 1
                dict.append(item)

            print("\rprogress：%.2f%%" % (float(cnt / total * 100)), end=' ')

        print(len(dict))
        delete = cnt - remain
        print('{} remain, {} loss, {} in total'.format(remain, delete, cnt))

    # 写入新json文件
    with open(new_json, 'a', encoding='utf-8') as write_f:
        write_f.write(json.dumps(dict, indent=4, ensure_ascii=False))

    print("write done")

if __name__ == '__main__':
    # parser = argparse.ArgumentParser()
    # parser.add_argument('image_dir', type=str,
    #                     help='image_dir')
    # parser.add_argument('grounding_json', type=str,
    #                     help='grounding_json')
    # parser.add_argument('new_json', type=str,
    #                     help='new_json')
    # args = parser.parse_args()
    #
    # image_dir = args.image_dir
    # grounding_json = args.grounding_json
    # new_json = args.new_json

    # python .\fit_grounding_json.py .\data\download .\data\grounding .\data\grounding\new.json
    image_dir = "./data/download" # 图片下载目录
    grounding_json = "./data/mm-event-graph/grounding" # grounding文件夹
    new_json = "./data/mm-event-graph/grounding/new.json" # 写入缺漏图片json


    fit_json(image_dir, grounding_json, new_json)



